from contextlib import asynccontextmanager

from typing import Literal

from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import AIMessage
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.prebuilt import ToolNode

from llm_utils import llm_qw3
from my_tools.log_utils import MyLogger
from sql_graph.node_prompt import generate_query_system_prompt, query_check_system
from sql_graph.sql_state import SQLState

# 引入日志
log = MyLogger().get_logger()

fastmcp_server_config = {
    "url": "http://127.0.0.1:6666/sse",
    "transport": "sse"
}


# 定义一个条件路由,第5个节点进行判断
def condition_route(state: SQLState) -> Literal[END, "n6_check_query"]:
    """条件路由，动态边"""
    message = state["messages"]
    last_message = message[-1]  # type: dict
    # 判断是否有工具调用，如果没有就返回 END
    if not last_message.tool_calls:
        return END
    return "n6_check_query"


# 数据库连接及使用langchain自带的数据库工具
db = SQLDatabase.from_uri('sqlite:///../chinook.db')
toolkit = SQLDatabaseToolkit(db=db, llm=llm_qw3)
tools = toolkit.get_tools()
# 获取表结构工具
get_schema_tool = next(t for t in tools if t.name == 'sql_db_schema')
# print(get_schema_tool.invoke('artists'))


"""第三个节点
当工作流执行到这个节点时，LLM 会分析之前的对话内容
如果需要了解数据库表结构，LLM 会自动调用 get_schema_tool
get_schema_tool 会查询数据库并返回表的结构信息
LLM 将这些信息整合到最终响应中
响应被返回到工作流中，供后续节点使用
"""
def n3_call_get_schema(state: SQLState):
    """ 第三个节点，执行调用工具的请求"""
    # FAQ: LangChain强制要求所有模型都接受 `tool_choice="any"`以及 `tool_choice=<工具名称字符串>` 这两种参数
    llm_with_tools = llm_qw3.bind_tools([get_schema_tool], tool_choice='any')
    response = llm_with_tools.invoke(state['messages'])
    return {"messages": [response]}


# 第四个节点: 直接使用langgraph提供的ToolNode
n4_get_schema_node = ToolNode([get_schema_tool], name="n4_get_schema")


@asynccontextmanager
async def make_graph():
    """定义，并编译工作流"""
    client = MultiServerMCPClient({'fw-mcp': fastmcp_server_config})
    tools = await client.get_tools()
    # 获取表名工具、执行sql工具
    list_tables_tool = next(t for t in tools if t.name == 'list_tables_tool')
    db_query_tool = next(tool for tool in tools if tool.name == "db_query_tool")
    # log.info(f'工具名称：{list_tables_tool}, 工具名称：{db_query_tool}')

    def n1_call_list_tables(state: SQLState):
        """第一个节点"""
        tool_call = {  #
            'name': 'list_tables_tool',
            'args': {},
            'id': 'abc123',
            'type': 'tool_call',  # 调用工具
        }
        tool_call_message = AIMessage(content='', tool_calls=[tool_call])
        return {"messages": [tool_call_message]}

    # 第二个节点，查询表名的工具节点
    n2_list_tables_tool = ToolNode(name='n2_list_tables_tool', tools=[list_tables_tool])

    # 第五个节点，用于生成sql语句
    def n5_generate_query(state: SQLState):
        """第五个节点,生成sql语句"""
        system_message = {
            'role': 'system',
            'content': generate_query_system_prompt,
        }
        # 这里不强带哦工具的调用，允许模型在获得解决方案时自然响应
        llm_with_tools = llm_qw3.bind_tools([db_query_tool])
        response = llm_with_tools.invoke([system_message] + state['messages'])
        return {'messages': [response]}

    # 第六个节点，执行sql检查的节点
    def n6_check_query(state: SQLState):
        """第六个节点，检查sql语句"""
        system_message = {
            'role': 'system',
            'content': query_check_system,
        }
        tool_call = state["messages"][-1].tool_calls[0]
        # 得到生成后的SQL
        user_message = {'role': 'user', 'content': tool_call["args"]["query"]}

        llm_with_tools = llm_qw3.bind_tools([db_query_tool], tool_choice='any')
        response = llm_with_tools.invoke([system_message, user_message])
        # If a message in `right` has the same ID as a message in `left`, the
        #    message from `right` will replace the message from `left`.
        response.id = state["messages"][-1].id
        return {'messages': [response]}

    # 第七个节点，执行sql语句的节点
    n7_run_query_node = ToolNode([db_query_tool], name='n7_run_query')

    """以下开始定义workflow,增加7个节点，一个条件路由"""
    workflow = StateGraph(SQLState)
    workflow.add_node(n1_call_list_tables)
    workflow.add_node(n2_list_tables_tool)
    workflow.add_node(n3_call_get_schema)
    workflow.add_node(n4_get_schema_node)
    workflow.add_node(n5_generate_query)
    workflow.add_node(n6_check_query)
    workflow.add_node(n7_run_query_node)

    workflow.add_edge(START, "n1_call_list_tables")
    workflow.add_edge("n1_call_list_tables", "n2_list_tables_tool")
    workflow.add_edge("n2_list_tables_tool", "n3_call_get_schema")
    workflow.add_edge("n3_call_get_schema", "n4_get_schema")
    workflow.add_edge("n4_get_schema", "n5_generate_query")
    workflow.add_conditional_edges('n5_generate_query', condition_route)
    workflow.add_edge('n6_check_query', 'n7_run_query')
    workflow.add_edge('n7_run_query', 'n5_generate_query')

    graph = workflow.compile()
    yield graph
